Cloud TPU の概要
Tensor Processing Unit（TPU）は、Google 独自に開発された特定用途向け集積回路（ASIC）であり、ML ワークロードの高速化に使用されます。TPU ハードウェアの詳細については、TPU アーキテクチャをご覧ください。Cloud TPU は、 Google Cloudで TPU をスケーラブルなコンピューティング リソースとして利用できるようにするウェブサービスです。
TPU は、ML アルゴリズムでよく見られる大規模な行列演算を実行するために設計されたハードウェアを使用して、モデルを効率的にトレーニングします。TPU にはオンチップの高帯域幅メモリ（HBM）が搭載されており、大きいモデルやバッチサイズを使用できます。TPU はスライスと呼ばれるグループに接続することで、コードをほとんど変更することなくワークロードをスケールアップできます。
TPU で実行されるコードは、アクセラレータ線形代数（XLA）コンパイラでコンパイルする必要があります。XLA は、ML フレームワーク アプリケーションによって出力されたグラフを受け取り、グラフの線形代数、損失、勾配のコンポーネントを TPU マシンコードにコンパイルするジャストインタイム コンパイラです。プログラムの残りの部分は TPU ホストマシンで実行されます。XLA コンパイラは、TPU ホストマシンで実行される TPU VM イメージの一部です。
Tensor Processing Unit の詳細については、TPU の考え方をご覧ください。
TPU の用途
Cloud TPU は特定のワークロード用に最適化されています。状況によっては、Compute Engine インスタンスで GPU または CPU を使用して ML ワークロードを実行する場合があります。一般に、次のガイドラインに基づいてワークロードに最適なハードウェアを決定できます。
CPU
- 最大限の柔軟性を必要とする迅速なプロトタイピング
- トレーニングに時間がかからない単純なモデル
- 実際のバッチサイズが小さい小規模なモデル
- C++ で記述されたカスタムの TensorFlow オペレーションが多く含まれるモデル
- ホストシステムの使用可能な I/O またはネットワーク帯域幅によって制限が課せられるモデル
GPU
- CPU 上で少なくとも部分的に実行しなければならない多数のカスタムの PyTorch / JAX 演算を使用するモデル
- Cloud TPU で利用できない TensorFlow 演算を使用するモデル（利用可能な TensorFlow 演算のリストをご覧ください）
- 実際のバッチサイズが大きい中～大規模なモデル
TPU
- 行列計算が多くを占めるモデル
- メインのトレーニング ループ内にカスタムの PyTorch / JAX 演算がないモデル
- トレーニングに数週間または数か月かかるモデル
- 実際のバッチサイズが大きい大規模なモデル
- 高度なランキングやレコメンデーションのワークロードでよくある極めて大規模なエンベディングを含むモデル
Cloud TPU は、以下のワークロードには適していません。
- 頻繁な分岐を必要とするか、要素ごとの代数演算を多く含む線形代数プログラム
- 高精度の演算を必要とするワークロード
- メインのトレーニング ループにカスタム演算が含まれるニューラル ネットワーク ワークロード
Google Cloudの TPU
TPU は、Cloud TPU VM、Google Kubernetes Engine、Vertex AI で使用できます。次の表に、各 Google Cloudサービスのリソースを示します。
| Google Cloud サービス | リソース |
|---|---|
| Cloud TPU | Cloud TPU VM を使ってみる |
| Google Kubernetes Engine | |
| Vertex AI |
モデル開発のベスト プラクティス
計算の多くを非行列演算（加算、再形成、連結など）が占めるプログラムは、MXU 使用率が高くならない可能性があります。次に、Cloud TPU に適したモデルを選択して作成する際に役立つガイドラインを示します。
レイアウト
XLA コンパイラでは、マトリックス ユニット（MXU）で効率的に計算を実行するために、行列乗算をより小さなブロックにタイリングするなどのコード変換が行われます。MXU ハードウェアの構造、128x128 のシストリック アレイ、8 の倍数の次元を基本とする TPU メモリ サブシステムの設計が、タイリングの効率化のために XLA コンパイラによって使用されます。
その結果、特定のレイアウトはタイリングが容易であるのに対し、他のレイアウトではタイリングの前に再形成が必要になります。再形成演算は多くの場合、Cloud TPU でメモリの制約を受けます。
形状
XLA コンパイラは、ジャストインタイムで最初のバッチの ML グラフをコンパイルします。後続のバッチの形状が異なる場合、モデルは機能しません（形状が変化するたびにグラフを再コンパイルするのでは遅すぎます）。したがって、テンソルの形状が動的なモデルは、TPU には適していません。
パディング
高パフォーマンスの Cloud TPU プログラムは、高密度なコンピューティングを 128x128 のチャンクにタイリングできるプログラムです。MXU 全体が行列計算で占有されない場合、コンパイラによってテンソルはゼロでパディングされます。パディングには次の 2 つの欠点があります。
- ゼロでパディングされたテンソルでは、TPU コアを十分に利用できません。
- パディングによって、テンソルに必要なオンチップ メモリ ストレージ量が増加し、極端な場合にはメモリ不足エラーにつながる可能性があります。
パディングは必要に応じて XLA コンパイラによって自動的に行われますが、op_profile ツールを使用して、行われるパディングの量を決定できます。TPU に適したテンソル次元を選択することで、パディングを回避できます。
次元
適切なテンソル次元を選択すると、TPU ハードウェア、特に MXU から最大のパフォーマンスを引き出すために役立ちます。XLA コンパイラは、バッチサイズまたは特徴量の次元のいずれかを使用して MXU を最大限に利用しようとします。したがって、これらのいずれかは 128 の倍数である必要があります。そうでない場合、コンパイラによっていずれかが 128 までパディングされます。理想的には、バッチサイズと特徴量の次元を 8 の倍数にすることで、メモリ サブシステムから高パフォーマンスを引き出すことができます。
Cloud TPU のスタートガイド
- Google Cloud アカウントを設定する
- Cloud TPU API を有効にする
- Cloud TPU に Cloud Storage バケットへのアクセスを許可する
- TPU で基本的な計算を実行する
- TPU でリファレンス モデルをトレーニングする
- モデルを分析する
サポートの依頼
サポートが必要な場合は、Cloud TPU サポートにお問い合わせください。有効な Google Cloud プロジェクトがある場合は、次の情報をご準備ください。
- Google Cloud プロジェクト ID
- TPU 名（存在する場合）
- ご提供いただく必要があるその他の情報
次のステップ
Cloud TPU の詳細を調べるには、次のリソースも役立ちます。